import numpy as np
import pandas as pd

# Preprocess data
data_train_fp = "./data/train.csv"
data_test_fp = './data/test.csv'


def load_data(file_path):
    data = pd.read_csv(file_path)
    data = data.drop(['PassengerId', 'Ticket', 'Cabin', 'Name'], axis=1)
    data['Sex'] = data['Sex'].map({'male': 0, 'female': 1})
    data['Embarked'] = data['Embarked'].map({'S': 0, 'C': 1, 'Q': 2})
    return data.fillna(data.mean())


def get_train_data(data_fp):
    data = load_data(data_fp)
    X = data.drop(['Survived'], axis=1).values
    y = data['Survived'].values
    return X, y


def get_test_data(data_fp):
    data = load_data(data_fp)
    return data.values


def sigmoid(z):
    z = np.clip(z, -709, 709)
    return 1 / (1 + np.exp(-z))


class SimpleLogicRegression:
    def __init__(self, learning_rate=0.01, epoch=100):
        self.weight = None
        self.bias = 0
        self.learning_rate = learning_rate
        self.epoch = epoch

    def fit(self, X, y):
        m, n = X.shape
        self.weight = np.zeros(n)

        for _ in range(self.epoch):
            model = np.dot(X, self.weight) + self.bias
            # print('epoch: %d, model: %s', self.epoch, model)
            y_pred = sigmoid(model)

            # 计算梯度
            dw = (1 / m) * np.dot(X.T, y_pred - y)
            db = (1 / m) * np.sum(y_pred - y)

            self.weight -= self.learning_rate * dw
            self.bias -= self.learning_rate * db

            # if self.epoch % 10 == 0:
            #     print("weight=", self.weight, ', bias', self.bias)

    def predict(self, X):
        model = np.dot(X, self.weight) + self.bias
        y_pred = sigmoid(model)
        y_pred_class = [1 if i > 0.5 else 0 for i in y_pred]
        return np.array(y_pred_class)


model = SimpleLogicRegression(learning_rate=0.01, epoch=10000)
X_train, y_train = get_train_data(data_train_fp)
model.fit(X_train, y_train)
print("Weight = ", model.weight)
print("Bias = ", model.bias)

X_test = get_test_data(data_test_fp)
y_test_pred = model.predict(X_test)
passenger_ids = pd.read_csv(data_test_fp)['PassengerId'].values
result_df = pd.DataFrame({
    'PassengerId': passenger_ids,
    'Survived': y_test_pred
})
# 将 DataFrame 写入到新的 CSV 文件
output_filepath = './data/gender_submission.csv'  # 替换为你想要的输出文件路径
result_df.to_csv(output_filepath, index=False)
